# 非定常的Stokes方程
from mpi4py import MPI
import numpy as np

import ufl
from basix.ufl import element
from dolfinx import default_real_type, fem, la
from dolfinx.fem import (
    Constant,
    Function,
    dirichletbc,
    form,
    functionspace,
)
from petsc4py import PETSc
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary
from ufl import div, dx, grad, inner, TestFunction, TrialFunction


# He Xiaoming 7.5.1 Example
def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data


def true_velocity(x, t):
    return [u_x(x, t), u_y(x, t)]


def true_pressure(x, t):
    return pressure(x, t)


def u_x(x, t):
    """
    x方向速度的真实解
    """
    return (x[0]**2 * x[1]**2 + np.exp(-x[1])) * np.cos(2 * np.pi * t)


def u_y(x, t):
    """
    y方向速度的真实解
    """
    return (-2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * t)


def u_init(x):
    """
    返回u: IC
    """
    t1 = (x[0]**2 * x[1]**2 + np.exp(-x[1]))
    t2 = -2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])
    return [t1, t2]


def p_init_expr(x):
    """
    返回p: IC
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1])


def p_end_expr(x):
    """
    返回p: IC
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * T)


def pressure(x, t):
    """
    压强的真实解
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * t)


# ufl格式
def f_x(x, t):
    """
    f的x分量
    """
    term1 = -2 * ufl.pi * (x[0]**2 * x[1]**2 +
                           ufl.exp(-x[1])) * ufl.sin(2 * ufl.pi * t)
    term2 = (
        -2 * nu * x[0]**2
        - 2 * nu * x[1]**2
        - nu * ufl.exp(-x[1])
        + ufl.pi**2 * ufl.cos(ufl.pi * x[0]) * ufl.cos(2 * ufl.pi * x[1])
    ) * ufl.cos(2 * ufl.pi * t)
    return term1 + term2


# He给出的数值算例不正确 第二项少了cos(2*pi*t)
def f_y(x, t):
    """
    f的y分量
    """
    term1 = -2 * ufl.pi * (-2/3 * x[0] * x[1]**3 + 2 - ufl.pi *
                           ufl.sin(ufl.pi * x[0])) * ufl.sin(2 * ufl.pi * t)
    term2 = (
        4 * nu * x[0] * x[1]
        - nu * ufl.pi**3 * ufl.sin(ufl.pi * x[0])
        + 2 * ufl.pi * (2 - ufl.pi * ufl.sin(ufl.pi *
                        x[0])) * ufl.sin(2 * ufl.pi * x[1])
    ) * ufl.cos(2 * ufl.pi * t)
    return term1 + term2


def f(x, t):
    """
    返回f(list)
    """
    return [f_x(x, t), f_y(x, t)]


def D(u):
    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)


def get_norm(t: float):
    """
    传入时间t
    """
    # 计算误差
    u_true = Function(V)  # 实例化空间
    p_true = Function(Q)

    u_true.interpolate(lambda x: np.array([u_x(x, t), u_y(x, t)]))
    def p_current(x):  # 只能这样设置时间 比较丑陋
        return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1]) * np.cos(2 * np.pi * t)
    p_true.interpolate(p_current)

    u_error = Function(V)
    u_error.x.array[:] = u_h.x.array - u_true.x.array
    p_error = Function(Q)
    p_error.x.array[:] = p_h.x.array - p_true.x.array

    # Compute norms of the u and p vectors
    u_comm = u_h.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    u_h1_error = form((grad(u_error)**2 + u_error**2) * ufl.dx)
    u_l2_error = form(u_error**2 * ufl.dx)
    u_h1_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_h1_error), MPI.SUM))
    u_l2_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_l2_error), MPI.SUM))  # 并行求和 返回float

    p_comm = p_h.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    p_h1_error = form((grad(p_error)**2 + p_error**2) * ufl.dx)
    p_l2_error = form(p_error**2 * ufl.dx)
    p_h1_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_h1_error), MPI.SUM))
    p_l2_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_l2_error), MPI.SUM))  # 并行求和 返回float

    # Linf 比较节点值 离散后区域无法计算理论最大值
    u_linf_E, p_linf_E = la.norm(
        u_error.x, la.Norm.linf), la.norm(p_error.x, la.Norm.linf)

    if MPI.COMM_WORLD.rank == 0:
        print(
            f"H1 norm for velocity: {u_h1_E}, pressure: {p_h1_E}.")
        print(
            f"L2 norm for velocity: {u_l2_E}, pressure: {p_l2_E}.")
        print(
            f"Linf norm for velocity: {u_linf_E}, pressure: {p_linf_E}.")


# Params
n_list = [
    [4, 1],
    [8, 2],
    [16, 4],
    [32, 8],
    [64, 16],
]
idx = int(input("请输入网格划分的索引: "))
nu = 1
t = 0  # 初始时刻
T = 1  # 终止时刻
time_steps = n_list[idx][0] ** 2  # 时间步 取得细致一点 否则误差无法减小...
dt = T / time_steps

msh = create_rectangle(
    MPI.COMM_WORLD, [np.array([0, -0.25]), np.array(
        [1, 0])], n=n_list[idx], cell_type=CellType.triangle
)
print(f"Mesh = {n_list[idx]}")

P2 = element(
    "Lagrange", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type
)
P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
V, Q = functionspace(msh, P2), functionspace(msh, P1)


# 标记边界
facets = locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(
        x[0], 1.0) | np.isclose(x[1], -0.25) | np.isclose(x[1], 0.0),
)
dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)


u_n = Function(V)
p_n = Function(Q)
u_n.interpolate(u_init)
p_n.interpolate(p_init_expr)
# 压力梯度是影响流体流动的关键 压力的绝对值并不会影响（物理角度）


for i in range(time_steps+1):
    print(f"t = {t}.")  # 循环的本质是不断求解Stokes方程 并将上一次的解作为这次的条件
    # Define trial and test functions
    u, p = TrialFunction(V), TrialFunction(Q)
    v, q = TestFunction(V), TestFunction(Q)
    # 预先计算边界条件的值
    bc_value = Function(V)
    bc_value.interpolate(lambda x: true_velocity(x, t))
    bc_v = dirichletbc(bc_value, dofs)

    # 这里x的定义不能放在循环外.. 原因未知
    x = ufl.SpatialCoordinate(msh)
    f_x_ = f_x(x, t)  # 确定时间t
    f_y_ = f_y(x, t)
    f_expr = ufl.as_vector([f_x_, f_y_])

    a_blocked = form(
        [
            [2 * nu * inner(D(u), D(v)) * dx + inner(u, v) /
             dt * dx, -inner(p, div(v)) * dx],
            [inner(div(u), q) * dx, None],
        ]
    )
    L_blocked = form(
        [inner(f_expr + u_n/dt, v) * dx,
         inner(Constant(msh, PETSc.ScalarType(0)), q) * dx]
    )
    # 预条件
    a_p11 = form(inner(p, q) * dx)
    a_p = [[a_blocked[0][0], None], [None, a_p11]]

    # Assemble Stokes problem
    A = assemble_matrix_block(a_blocked, bcs=[bc_v])
    A.assemble()
    P = assemble_matrix_block(a_p, bcs=[bc_v])
    P.assemble()
    b = assemble_vector_block(L_blocked, a_blocked, bcs=[bc_v])
    
    # Set the nullspace for pressure (since pressure is determined only
    # up to a constant)
    null_vec = A.createVecLeft()
    offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs

    null_vec.array[offset:] = 1.0
    null_vec.normalize()
    nsp = PETSc.NullSpace().create(vectors=[null_vec])
    assert nsp.test(A)
    A.setNullSpace(nsp)
    
    # 这里条件数很大!!! 数量级是10^17
    # stiff_mat = getStiffnessMatrix(A=A)  # 刚度矩阵太大会oom
    # load_vec = b.getArray()  # 载荷向量

    # Create and configure solver
    ksp = PETSc.KSP().create(msh.comm)  # type: ignore
    ksp.setOperators(A)
    ksp.setType("preonly")

    pc = ksp.getPC()
    pc.setType("lu")
    sys = PETSc.Sys()  # type: ignore
    use_superlu = PETSc.IntType == np.int64
    if sys.hasExternalPackage("mumps") and not use_superlu:
        pc.setFactorSolverType("mumps")
        pc.setFactorSetUpSolverType()
        # icntl=24 允许处理奇异矩阵 =25禁用错误检测
        pc.getFactorMatrix().setMumpsIcntl(icntl=24, ival=1)
        pc.getFactorMatrix().setMumpsIcntl(icntl=25, ival=0)
    else:
        pc.setFactorSolverType("superlu_dist")

    # Solve
    x = A.createVecRight()
    ksp.solve(b, x)

    # Split the solution
    u_h = fem.Function(V)  # 存储结果
    p_h = fem.Function(Q)
    offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
    u_h.x.array[:offset] = x.array_r[:offset]
    u_h.x.scatter_forward()
    p_h.x.array[: (len(x.array_r) - offset)] = x.array_r[offset:]
    p_h.x.scatter_forward()

    # 迭代求解下一时刻的数值解
    

    # 更新u_n和p_n的解
    # 注释掉这行突然就能跑了 因为初始时刻和终止时刻边界条件正好一样
    u_n.x.array[:] = u_h.x.array  # 从当前解中提取速度
    p_n.x.array[:] = p_h.x.array  # 从当前解中提取压力
    u_n.x.scatter_forward()
    p_n.x.scatter_forward()
    
    t += dt
get_norm(t=T)  # 计算误差
